// Usage: node flatten_switch_unroll.js [inputPath] [outputPath]
// Defaults: inputPath = ./deobfuscated_code.js, outputPath = ./deobfuscated_code.unflat.js

const fs = require("fs");
const path = require("path");
const parser = require("@babel/parser");
const traverse = require("@babel/traverse").default;
const t = require("@babel/types");
const generate = require("@babel/generator").default;

const inputPath = process.argv[2] || path.resolve("./de_de_de_code2.js");
const outputPath = process.argv[3] || path.resolve("./de_de_de_de_code2.js");

const code = fs.readFileSync(inputPath, "utf8");
const ast = parser.parse(code, {
  sourceType: "unambiguous",
  plugins: [
    "optionalChaining",
    "nullishCoalescingOperator",
    "objectRestSpread",
    "classProperties",
    "classPrivateProperties",
    "classPrivateMethods",
    "bigInt",
    "numericSeparator",
    "logicalAssignment",
    "topLevelAwait"
  ]
});

function getSplitArray(node) {
  // matches: "0|5|4|3|2|1"["split"]('|') OR "0|1|2".split('|')
  if (!t.isCallExpression(node)) return null;
  const { callee, arguments: args } = node;
  if (args.length !== 1) return null;
  const isSplit = (
    t.isMemberExpression(callee) &&
    ((t.isIdentifier(callee.property, { name: "split" }) && !callee.computed) ||
     (t.isStringLiteral(callee.property) && callee.property.value === "split"))
  );
  if (!isSplit) return null;
  const obj = callee.object;
  if (!t.isStringLiteral(obj)) return null;
  const sep = args[0];
  if (!(t.isStringLiteral(sep) && sep.value.length >= 1)) return null;
  return obj.value.split(sep.value);
}

function isAlwaysTrue(node) {
  if (!node) return false;
  if (t.isBooleanLiteral(node, { value: true })) return true;
  // !![]
  if (
    t.isUnaryExpression(node, { operator: "!" }) &&
    t.isUnaryExpression(node.argument, { operator: "!" }) &&
    t.isArrayExpression(node.argument.argument)
  ) return true;
  return false;
}

function isIndexingProgress(node, idxName, arrName) {
  // matches s[W++] or s[W+=1] etc; we expect s[W++] exactly in most cases
  if (!t.isMemberExpression(node)) return false;
  if (!node.computed) return false;
  if (!t.isIdentifier(node.object, { name: arrName })) return false;
  const p = node.property;
  if (t.isUpdateExpression(p) && p.operator === "++" && !p.prefix && t.isIdentifier(p.argument, { name: idxName })) return true;
  return false;
}

function stripTrailingContinue(stmts) {
  const out = [];
  for (let i = 0; i < stmts.length; i++) {
    const s = stmts[i];
    if (t.isContinueStatement(s)) continue; // drop case-local continue
    out.push(s);
  }
  return out;
}

let transformed = 0;

traverse(ast, {
  ForStatement(path) {
    const node = path.node;
    if (!isAlwaysTrue(node.test)) return;

    // try to detect: for (s = "...".split('|'), W = 0; !![];) { switch (s[W++]) { ... } break; }

    // find s & order from init (SequenceExpression or comma-separated assignments)
    let arrName = null, idxName = null, order = null;

    const tryInit = (init) => {
      if (!init) return;
      const assignments = [];
      if (t.isSequenceExpression(init)) {
        for (const ex of init.expressions) assignments.push(ex);
      } else if (t.isAssignmentExpression(init)) {
        assignments.push(init);
      } else if (t.isVariableDeclaration(init)) {
        for (const d of init.declarations) if (d.init) assignments.push(t.assignmentExpression("=", d.id, d.init));
      }
      for (const a of assignments) {
        if (!t.isAssignmentExpression(a, { operator: "=" })) continue;
        if (!t.isIdentifier(a.left)) continue;
        // s = "...".split('|')
        const arr = getSplitArray(a.right);
        if (arr) {
          arrName = a.left.name;
          order = arr;
          continue;
        }
        // W = 0
        if (t.isNumericLiteral(a.right) && a.right.value === 0 && t.isIdentifier(a.left)) {
          idxName = a.left.name;
        }
      }
    };

    tryInit(node.init);
    if (!arrName || !idxName || !order) return; // not our pattern

    // body must contain a SwitchStatement on s[W++] and then a BreakStatement (to exit once all cases run)
    if (!t.isBlockStatement(node.body) || node.body.body.length < 1) return;
    const first = node.body.body[0];
    if (!t.isSwitchStatement(first)) return;
    if (!isIndexingProgress(first.discriminant, idxName, arrName)) return;

    // collect case bodies by label
    const table = new Map();
    for (const cs of first.cases) {
      if (!cs.test || !(t.isStringLiteral(cs.test) || t.isNumericLiteral(cs.test))) continue;
      const key = String(t.isStringLiteral(cs.test) ? cs.test.value : cs.test.value);
      table.set(key, stripTrailingContinue(cs.consequent));
    }

    // flatten in given order
    const flatStatements = [];
    for (const key of order) {
      const block = table.get(String(key));
      if (!block) continue; // tolerate missing
      for (const st of block) flatStatements.push(t.cloneNode(st, true));
    }

    // Replace entire for-loop with the flattened sequence
    if (flatStatements.length > 0) {
      path.replaceWith(t.blockStatement(flatStatements));
      transformed++;
    }
  }
});

const output = generate(ast, { compact: false, comments: true }).code;
fs.writeFileSync(outputPath, output, "utf8");
console.log(`[OK] Flattened blocks: ${transformed}`);
console.log(`[OUT] ${outputPath}`);
